import os
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import io

# 初始化 InferenceClient
def init_client():
    """初始化 HuggingFace InferenceClient"""
    try:
        client = InferenceClient(
            provider="fal-ai",
            api_key=os.environ["HF_API_KEY"],
        )
        return client
    except KeyError:
        raise gr.Error("请设置环境变量 HF_API_KEY")

def compress_image(image, max_size_mb=3.5, quality=85):
    """
    压缩图像以确保不超过最大文件大小
    
    Args:
        image: PIL Image 对象
        max_size_mb: 最大文件大小（MB），默认 3.5MB 留出安全余量
        quality: JPEG 质量（1-100）
    
    Returns:
        压缩后的 PIL Image 对象
    """
    max_size_bytes = max_size_mb * 1024 * 1024
    
    # 首先尝试调整图像尺寸
    width, height = image.size
    scale_factor = 1.0
    
    # 如果图像太大，按比例缩小
    if width > 1024 or height > 1024:
        scale_factor = min(1024 / width, 1024 / height)
        new_width = int(width * scale_factor)
        new_height = int(height * scale_factor)
        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 转换为 RGB 模式（如果不是的话）
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # 尝试不同的质量设置
    for q in range(quality, 30, -10):
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='JPEG', quality=q, optimize=True)
        img_byte_arr.seek(0)
        
        if img_byte_arr.tell() <= max_size_bytes:
            img_byte_arr.seek(0)
            return Image.open(img_byte_arr)
    
    # 如果还是太大，进一步缩小尺寸
    while scale_factor > 0.3:
        scale_factor *= 0.8
        new_width = int(width * scale_factor)
        new_height = int(height * scale_factor)
        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        img_byte_arr = io.BytesIO()
        resized_image.save(img_byte_arr, format='JPEG', quality=60, optimize=True)
        img_byte_arr.seek(0)
        
        if img_byte_arr.tell() <= max_size_bytes:
            img_byte_arr.seek(0)
            return Image.open(img_byte_arr)
    
    raise gr.Error("无法将图像压缩到合适的大小，请尝试使用更小的图像")

def process_image_to_image(input_image, prompt, model_name="black-forest-labs/FLUX.1-Kontext-dev"):
    """
    使用 FLUX.1 模型进行图像到图像转换
    
    Args:
        input_image: 输入的 PIL Image 对象
        prompt: 文本提示词
        model_name: 模型名称
    
    Returns:
        转换后的图像
    """
    try:
        # 初始化客户端
        client = init_client()
        
        # 压缩图像
        compressed_image = compress_image(input_image)
        
        # 将 PIL Image 转换为字节
        img_byte_arr = io.BytesIO()
        compressed_image.save(img_byte_arr, format='JPEG', quality=85, optimize=True)
        img_byte_arr = img_byte_arr.getvalue()
        
        # 检查文件大小
        file_size_mb = len(img_byte_arr) / (1024 * 1024)
        print(f"图像文件大小: {file_size_mb:.2f} MB")
        
        if file_size_mb > 4:
            raise gr.Error(f"图像文件仍然过大 ({file_size_mb:.2f} MB)，请使用更小的图像")
        
        # 调用模型进行图像转换
        output_image = client.image_to_image(
            img_byte_arr,
            prompt=prompt,
            model=model_name,
        )
        
        return output_image
        
    except Exception as e:
        if "413" in str(e) or "Payload Too Large" in str(e):
            raise gr.Error("图像文件过大，请使用更小的图像或尝试不同的图像")
        else:
            raise gr.Error(f"图像处理失败: {str(e)}")

def create_gradio_interface():
    """创建 Gradio 界面"""
    
    # 定义界面
    with gr.Blocks(title="FLUX.1 图像转换工具", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🎨 基于FLUX模型与Gradio框架的AI图像转换工具
        
        使用 FLUX.1 模型将您的图像转换为新的风格。上传一张图片，输入提示词，然后生成全新的图像！
        作者：RoysterCDD
        
        ### 使用说明：
        1. 上传一张图片（建议小于 4MB）
        2. 输入描述性的提示词（例如："Turn the cat into a tiger"）
        3. 点击"生成图像"按钮
        4. 等待处理完成，查看结果
        
        **注意：** 如果图像过大，系统会自动压缩图像以确保处理成功。
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # 输入区域
                gr.Markdown("### 📤 输入")
                
                input_image = gr.Image(
                    label="上传图片（建议小于 4MB）",
                    type="pil",
                    height=300
                )
                
                prompt = gr.Textbox(
                    label="提示词",
                    placeholder="例如：Turn the cat into a tiger, Make it more colorful...",
                    lines=3,
                    max_lines=5
                )
                
                model_selector = gr.Dropdown(
                    choices=["black-forest-labs/FLUX.1-Kontext-dev"],
                    value="black-forest-labs/FLUX.1-Kontext-dev",
                    label="选择模型",
                    interactive=False
                )
                
                generate_btn = gr.Button(
                    "🎨 生成图像",
                    variant="primary",
                    size="lg"
                )
                
                # 状态显示
                status = gr.Textbox(
                    label="状态",
                    value="准备就绪",
                    interactive=False
                )
            
            with gr.Column(scale=1):
                # 输出区域
                gr.Markdown("### 📥 输出")
                
                output_image = gr.Image(
                    label="生成结果",
                    height=400
                )
                
                # 下载按钮（移除，因为图像组件自带下载功能）
                pass
        
        # 示例区域
        with gr.Accordion("💡 示例提示词", open=False):
            gr.Markdown("""
            **动物转换：**
            - Turn the cat into a tiger
            - Transform the dog into a wolf
            - Make the bird look like an eagle
            
            **风格转换：**
            - Make it more colorful and vibrant
            - Convert to watercolor painting style
            - Transform into anime style
            
            **环境转换：**
            - Put the subject in a forest
            - Place it in a futuristic city
            - Set it in a magical garden
            """)
        
        # 处理函数
        def generate_image(img, text, model):
            if img is None:
                raise gr.Error("请先上传一张图片")
            if not text.strip():
                raise gr.Error("请输入提示词")
            
            try:
                result = process_image_to_image(img, text, model)
                return result, "图像生成完成！"
            except Exception as e:
                raise gr.Error(f"图像处理失败: {str(e)}")
        
        # 绑定事件
        generate_btn.click(
            fn=generate_image,
            inputs=[input_image, prompt, model_selector],
            outputs=[output_image, status],
            api_name="generate"
        )
        
        # 移除下载按钮相关代码，因为图像组件自带下载功能
        pass
    
    return demo

if __name__ == "__main__":
    # 创建并启动 Gradio 应用
    app = create_gradio_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        debug=True
    )
